#!/usr/bin/env python3
"""
数据下载示例

演示如何使用增强数据下载器批量下载历史数据
"""

import sys
import os
from datetime import datetime, timedelta
from pathlib import Path

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from xcquant.data.data_manager import DataManager
from xcquant.data.enhanced_data_downloader import EnhancedDataDownloader, DownloadTask, create_common_download_tasks
from xcquant.core.unified_core import ConfigManager
from xcquant.core import unified_logger


def download_basic_data():
    """下载基础数据示例"""
    
    logger = unified_logger.get_logger('data_download_basic')
    logger.info("开始下载基础数据")
    
    try:
        # 1. 创建数据管理器
        data_config = {
            'data_path': './data',
            'cache_size': 1000,
            'auto_download': True,
            'data_sources': ['okx', 'binance']
        }
        
        data_manager = DataManager(data_config)
        
        # 2. 创建增强数据下载器
        download_config = {
            'max_workers': 2,           # 并发数量
            'chunk_size_days': 7,       # 每块7天数据
            'retry_attempts': 3,        # 重试次数
            'retry_delay': 1.0,         # 重试延迟
            'rate_limit_delay': 0.2     # 速率限制延迟
        }
        
        downloader = EnhancedDataDownloader(data_manager, download_config)
        
        # 3. 创建下载任务
        end_date = datetime.now()
        start_date = end_date - timedelta(days=30)  # 下载最近30天数据
        
        # 主要交易对
        symbols = ['BTC-USDT', 'ETH-USDT', 'BNB-USDT']
        intervals = ['1h', '4h', '1d']
        exchanges = ['okx']
        
        tasks = []
        for exchange in exchanges:
            for symbol in symbols:
                for interval in intervals:
                    # 现货任务
                    task_spot = DownloadTask(
                        exchange=exchange,
                        symbol=symbol,
                        interval=interval,
                        trade_type='spot',
                        start_date=start_date,
                        end_date=end_date,
                        priority=1
                    )
                    tasks.append(task_spot)
                    
                    # 合约任务
                    task_futures = DownloadTask(
                        exchange=exchange,
                        symbol=symbol,
                        interval=interval,
                        trade_type='futures',
                        start_date=start_date,
                        end_date=end_date,
                        priority=1
                    )
                    tasks.append(task_futures)
        
        # 4. 添加任务到下载器
        task_ids = downloader.add_batch_tasks(tasks)
        logger.info(f"添加了 {len(task_ids)} 个下载任务")
        
        # 5. 定义进度回调
        def progress_callback(progress):
            logger.info(f"下载进度: {progress['progress_percent']:.1f}% "
                       f"({progress['completed_tasks']}/{progress['total_tasks']})")
        
        def completion_callback(stats):
            logger.info("下载完成!")
            logger.info(f"总任务: {stats['total_tasks']}")
            logger.info(f"成功: {stats['completed_tasks']}")
            logger.info(f"失败: {stats['failed_tasks']}")
            logger.info(f"数据点: {stats['total_data_points']}")
            
            if stats['start_time'] and stats['end_time']:
                duration = (stats['end_time'] - stats['start_time']).total_seconds()
                logger.info(f"耗时: {duration:.1f} 秒")
        
        # 6. 开始下载
        if downloader.start_download(progress_callback, completion_callback):
            logger.info("下载已启动，等待完成...")
            
            # 等待下载完成
            import time
            while downloader.is_downloading:
                time.sleep(2)
                
                # 显示队列状态
                status = downloader.get_queue_status()
                if status['active_tasks'] > 0:
                    logger.info(f"活跃任务: {status['active_tasks']}, "
                               f"已完成: {status['completed_tasks']}")
            
            # 7. 生成下载报告
            report_path = downloader.export_progress_report('./download_report.json')
            if report_path:
                logger.info(f"下载报告已生成: {report_path}")
            
            logger.info("基础数据下载完成")
            return True
        else:
            logger.error("启动下载失败")
            return False
            
    except Exception as e:
        logger.error(f"下载基础数据失败: {e}")
        return False


def download_comprehensive_data():
    """下载全面数据示例"""
    
    logger = unified_logger.get_logger('data_download_comprehensive')
    logger.info("开始下载全面数据")
    
    try:
        # 1. 创建数据管理器
        data_manager = DataManager({
            'data_path': './comprehensive_data',
            'cache_size': 5000,
            'auto_download': True
        })
        
        # 2. 创建下载器 (更高并发)
        downloader = EnhancedDataDownloader(data_manager, {
            'max_workers': 4,
            'chunk_size_days': 5,
            'retry_attempts': 5,
            'rate_limit_delay': 0.1
        })
        
        # 3. 使用便捷函数创建常用任务
        logger.info("创建常用下载任务...")
        
        # 主流币种
        main_symbols = [
            'BTC-USDT', 'ETH-USDT', 'BNB-USDT', 'ADA-USDT', 'SOL-USDT',
            'XRP-USDT', 'DOT-USDT', 'DOGE-USDT', 'AVAX-USDT', 'MATIC-USDT'
        ]
        
        # 全时间周期
        all_intervals = ['1m', '5m', '15m', '30m', '1h', '4h', '1d']
        
        # 多交易所
        exchanges = ['okx', 'binance']
        
        tasks = create_common_download_tasks(
            exchanges=exchanges,
            symbols=main_symbols,
            intervals=all_intervals,
            days_back=90  # 下载3个月数据
        )
        
        logger.info(f"创建了 {len(tasks)} 个下载任务")
        
        # 4. 分批添加任务 (避免内存占用过大)
        batch_size = 50
        total_batches = (len(tasks) + batch_size - 1) // batch_size
        
        for i in range(0, len(tasks), batch_size):
            batch = tasks[i:i + batch_size]
            batch_num = i // batch_size + 1
            
            logger.info(f"处理批次 {batch_num}/{total_batches} ({len(batch)} 个任务)")
            
            # 添加批次任务
            downloader.add_batch_tasks(batch)
            
            # 启动下载
            def batch_progress_callback(progress):
                if progress['progress_percent'] % 10 == 0:  # 每10%显示一次
                    logger.info(f"批次 {batch_num} 进度: {progress['progress_percent']:.1f}%")
            
            def batch_completion_callback(stats):
                logger.info(f"批次 {batch_num} 完成: 成功 {stats['completed_tasks']}, "
                           f"失败 {stats['failed_tasks']}")
            
            if downloader.start_download(batch_progress_callback, batch_completion_callback):
                # 等待当前批次完成
                import time
                while downloader.is_downloading:
                    time.sleep(5)
                
                # 清空队列准备下一批次
                downloader.clear_queue()
            
            logger.info(f"批次 {batch_num} 处理完成")
        
        logger.info("全面数据下载完成")
        return True
        
    except Exception as e:
        logger.error(f"下载全面数据失败: {e}")
        return False


def download_custom_data():
    """自定义数据下载示例"""
    
    logger = unified_logger.get_logger('data_download_custom')
    logger.info("开始自定义数据下载示例")
    
    try:
        # 获取用户输入
        print("自定义数据下载配置")
        print("=" * 30)
        
        # 交易所选择
        print("支持的交易所: okx, binance")
        exchange = input("请输入交易所 (默认: okx): ").strip() or 'okx'
        
        # 交易对选择
        print("示例交易对: BTC-USDT, ETH-USDT, BNB-USDT")
        symbols_input = input("请输入交易对 (多个用逗号分隔): ").strip()
        symbols = [s.strip().upper() for s in symbols_input.split(',') if s.strip()]
        
        if not symbols:
            symbols = ['BTC-USDT']
        
        # 时间周期选择
        print("支持的时间周期: 1m, 5m, 15m, 30m, 1h, 4h, 1d")
        intervals_input = input("请输入时间周期 (多个用逗号分隔, 默认: 1h,4h,1d): ").strip()
        if intervals_input:
            intervals = [s.strip() for s in intervals_input.split(',') if s.strip()]
        else:
            intervals = ['1h', '4h', '1d']
        
        # 交易类型选择
        trade_types_input = input("请输入交易类型 (spot,futures 默认: spot,futures): ").strip()
        if trade_types_input:
            trade_types = [s.strip() for s in trade_types_input.split(',') if s.strip()]
        else:
            trade_types = ['spot', 'futures']
        
        # 时间范围选择
        days_input = input("请输入下载天数 (默认: 30): ").strip()
        try:
            days_back = int(days_input) if days_input else 30
        except ValueError:
            days_back = 30
        
        # 创建数据管理器和下载器
        data_manager = DataManager({'data_path': './custom_data'})
        downloader = EnhancedDataDownloader(data_manager, {
            'max_workers': 3,
            'retry_attempts': 3
        })
        
        # 创建任务
        end_date = datetime.now()
        start_date = end_date - timedelta(days=days_back)
        
        tasks = []
        for symbol in symbols:
            for interval in intervals:
                for trade_type in trade_types:
                    task = DownloadTask(
                        exchange=exchange,
                        symbol=symbol,
                        interval=interval,
                        trade_type=trade_type,
                        start_date=start_date,
                        end_date=end_date,
                        priority=1
                    )
                    tasks.append(task)
        
        logger.info(f"创建了 {len(tasks)} 个自定义下载任务")
        logger.info(f"交易所: {exchange}")
        logger.info(f"交易对: {symbols}")
        logger.info(f"时间周期: {intervals}")
        logger.info(f"交易类型: {trade_types}")
        logger.info(f"时间范围: {start_date.strftime('%Y-%m-%d')} 到 {end_date.strftime('%Y-%m-%d')}")
        
        # 确认下载
        confirm = input("\n确认开始下载? (y/N): ").strip().lower()
        if confirm != 'y':
            logger.info("下载已取消")
            return False
        
        # 添加任务并开始下载
        downloader.add_batch_tasks(tasks)
        
        def progress_callback(progress):
            print(f"\r下载进度: {progress['progress_percent']:.1f}% "
                  f"({progress['completed_tasks']}/{progress['total_tasks']})", end='', flush=True)
        
        def completion_callback(stats):
            print(f"\n下载完成! 成功: {stats['completed_tasks']}, 失败: {stats['failed_tasks']}")
        
        if downloader.start_download(progress_callback, completion_callback):
            # 等待完成
            import time
            while downloader.is_downloading:
                time.sleep(1)
            
            logger.info("自定义数据下载完成")
            return True
        else:
            logger.error("启动下载失败")
            return False
            
    except Exception as e:
        logger.error(f"自定义数据下载失败: {e}")
        return False


def check_data_status():
    """检查数据状态"""
    
    logger = unified_logger.get_logger('data_status')
    logger.info("检查数据状态")
    
    try:
        data_manager = DataManager({'data_path': './data'})
        
        # 获取数据统计
        stats = data_manager.get_data_stats()
        
        print("\n数据状态报告")
        print("=" * 50)
        print(f"数据目录: {stats.get('data_path', 'N/A')}")
        print(f"总文件数: {stats.get('total_files', 0)}")
        print(f"总数据量: {stats.get('total_size_mb', 0):.2f} MB")
        print(f"缓存大小: {stats.get('cache_size', 0)}")
        print(f"缓存命中率: {stats.get('cache_hit_rate', 0):.1%}")
        
        # 按交易所统计
        exchange_stats = stats.get('by_exchange', {})
        if exchange_stats:
            print("\n按交易所统计:")
            for exchange, count in exchange_stats.items():
                print(f"  {exchange}: {count} 个文件")
        
        # 按交易对统计
        symbol_stats = stats.get('by_symbol', {})
        if symbol_stats:
            print("\n按交易对统计 (前10个):")
            sorted_symbols = sorted(symbol_stats.items(), key=lambda x: x[1], reverse=True)[:10]
            for symbol, count in sorted_symbols:
                print(f"  {symbol}: {count} 个文件")
        
        # 按时间周期统计
        interval_stats = stats.get('by_interval', {})
        if interval_stats:
            print("\n按时间周期统计:")
            for interval, count in interval_stats.items():
                print(f"  {interval}: {count} 个文件")
        
        print("\n数据状态检查完成")
        return True
        
    except Exception as e:
        logger.error(f"检查数据状态失败: {e}")
        return False


def main():
    """主函数"""
    print("XCQuant 数据下载工具")
    print("=" * 50)
    print("1. 下载基础数据 (主要币种, 30天)")
    print("2. 下载全面数据 (多币种, 多周期, 90天)")
    print("3. 自定义数据下载")
    print("4. 检查数据状态")
    print("5. 退出")
    
    while True:
        try:
            choice = input("\n请选择操作 (1-5): ").strip()
            
            if choice == '1':
                print("\n开始下载基础数据...")
                success = download_basic_data()
                if success:
                    print("✅ 基础数据下载成功")
                else:
                    print("❌ 基础数据下载失败")
                    
            elif choice == '2':
                print("\n开始下载全面数据...")
                success = download_comprehensive_data()
                if success:
                    print("✅ 全面数据下载成功")
                else:
                    print("❌ 全面数据下载失败")
                    
            elif choice == '3':
                print("\n开始自定义数据下载...")
                success = download_custom_data()
                if success:
                    print("✅ 自定义数据下载成功")
                else:
                    print("❌ 自定义数据下载失败")
                    
            elif choice == '4':
                check_data_status()
                    
            elif choice == '5':
                print("退出程序")
                break
                
            else:
                print("无效选择，请输入 1-5")
                
        except KeyboardInterrupt:
            print("\n\n程序被用户中断")
            break
        except Exception as e:
            print(f"运行错误: {e}")


if __name__ == '__main__':
    main()